import backtrader as bt
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import akshare as ak
import warnings
from datetime import datetime, timedelta
import time
from tqdm import tqdm

warnings.filterwarnings('ignore')

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


class StockScanner:
    """股票扫描器：自动筛选满足技术条件的股票"""

    def __init__(self):
        self.condition_stocks = []

    def get_all_stocks(self):
        """获取股票列表"""
        try:
            stock_list = ak.stock_info_a_code_name()
            if stock_list is not None and not stock_list.empty:
                print(f"成功获取 {len(stock_list)} 只股票")
                return stock_list
            else:
                print("从akshare获取股票列表失败，使用预设股票列表")
        except Exception as e:
            print(f"获取股票列表时出错: {e}")

        # 使用预设的股票列表
        common_stocks = [
            ('000001', '平安银行'), ('600036', '招商银行'), ('601318', '中国平安'),
            ('600519', '贵州茅台'), ('000858', '五粮液'), ('600276', '恒瑞医药'),
            ('000333', '美的集团'), ('000651', '格力电器'), ('600887', '伊利股份'),
            ('600900', '长江电力'), ('601166', '兴业银行'), ('601288', '农业银行'),
            ('601398', '工商银行'), ('601988', '中国银行'), ('601857', '中国石油'),
            ('600000', '浦发银行'), ('600016', '民生银行'), ('600030', '中信证券'),
            ('600048', '保利地产'), ('600104', '上汽集团'), ('600309', '万华化学'),
            ('600585', '海螺水泥'), ('600837', '海通证券'), ('601088', '中国神华'),
            ('601169', '北京银行'), ('601186', '中国铁建'), ('601211', '国泰君安'),
            ('601328', '交通银行'), ('601601', '中国太保'), ('601628', '中国人寿'),
            ('601668', '中国建筑'), ('601688', '华泰证券'), ('601818', '光大银行'),
            ('601888', '中国中免'), ('601989', '中国重工')
        ]
        return pd.DataFrame(common_stocks, columns=['code', 'name'])

    def calculate_technical_indicators(self, df):
        """计算技术指标"""
        if df is None or len(df) < 250:  # 只需要250日数据
            return None

        try:
            # 计算移动平均线
            ma5 = df['close'].rolling(window=5).mean()
            ma10 = df['close'].rolling(window=10).mean()
            ma89 = df['close'].rolling(window=89).mean()
            ma120 = df['close'].rolling(window=120).mean()
            ma181 = df['close'].rolling(window=181).mean()
            ma250 = df['close'].rolling(window=250).mean()

            # 计算14日指数移动平均线
            ema14 = df['close'].ewm(span=14, adjust=False).mean()

            # 检查是否有足够的数据
            if len(ma5) < 2 or len(ema14) < 1:
                return None

            # 1. 均线多头排列：5日>10日>89日>120日>181日>250日
            ma_bullish = (ma5.iloc[-1] > ma10.iloc[-1] and
                          ma10.iloc[-1] > ma89.iloc[-1] and
                          ma89.iloc[-1] > ma120.iloc[-1] and
                          ma120.iloc[-1] > ma181.iloc[-1] and
                          ma181.iloc[-1] > ma250.iloc[-1] and
                          ma5.iloc[-1] > ma5.iloc[-2])

            # 2. K线沿14日EMA且不破
            current_close = df['close'].iloc[-1]
            current_ema = ema14.iloc[-1]

            # K线在EMA上方且距离适中（不超过5%）
            above_ema = current_close > current_ema
            ema_distance = abs(current_close - current_ema) / current_ema if current_ema > 0 else 0
            near_ema = ema_distance < 0.05

            # 检查近期是否跌破EMA（过去5天）
            recent_below_ema = False
            for i in range(1, min(6, len(df))):
                if df['close'].iloc[-i] < ema14.iloc[-i]:
                    recent_below_ema = True
                    break

            ema_condition = above_ema and near_ema and not recent_below_ema

            return {
                'ma5': ma5.iloc[-1],
                'ma10': ma10.iloc[-1],
                'ma89': ma89.iloc[-1],
                'ma120': ma120.iloc[-1],
                'ma181': ma181.iloc[-1],
                'ma250': ma250.iloc[-1],
                'ema14': ema14.iloc[-1],
                'close': current_close,
                'ma_bullish': ma_bullish,
                'ema_condition': ema_condition,
                'ema_distance': ema_distance * 100
            }

        except Exception as e:
            print(f"计算技术指标时出错: {e}")
            return None

    def is_strategy_condition(self, indicators):
        """判断是否满足策略条件"""
        if indicators is None:
            return False

        # 只保留均线多头排列和EMA条件
        conditions = [
            indicators['ma_bullish'],  # 均线多头排列
            indicators['ema_condition'],  # K线沿EMA且不破
        ]

        return all(conditions)

    def scan_stocks(self, start_date='20230101', end_date='20231231', max_stocks=10000):
        """扫描满足条件的股票"""
        print("开始扫描满足策略条件的股票...")
        print("策略条件: 5/10/89/120/181/250日均线多头排列 + K线沿14日EMA不破")

        stock_list = self.get_all_stocks()
        if stock_list is None or stock_list.empty:
            print("无法获取股票列表，请检查网络连接或akshare配置")
            return []

        selected_stocks = []
        total_stocks = min(5000, len(stock_list))

        print(f"总共扫描 {total_stocks} 只股票")

        for idx, row in stock_list.head(total_stocks).iterrows():
            if len(selected_stocks) >= max_stocks:
                break

            symbol = str(row['code']).zfill(6)
            name = row['name']

            try:
                # 获取股票数据
                df = ak.stock_zh_a_hist(symbol=symbol, period="daily",
                                        start_date=start_date, end_date=end_date,
                                        adjust="qfq")

                if df is None or df.empty or len(df) < 250:
                    continue

                # 处理数据
                df = df.rename(columns={
                    '日期': 'date', '开盘': 'open', '最高': 'high', '最低': 'low',
                    '收盘': 'close', '成交量': 'volume'
                })
                df['date'] = pd.to_datetime(df['date'])
                df.set_index('date', inplace=True)
                df = df[['open', 'high', 'low', 'close', 'volume']]
                df = df.astype(float)

                # 计算技术指标
                indicators = self.calculate_technical_indicators(df)
                if indicators is None:
                    continue

                # 检查是否满足策略条件
                if self.is_strategy_condition(indicators):
                    selected_stocks.append({
                        'symbol': symbol,
                        'name': name,
                        'price': indicators['close'],
                        'ma5': indicators['ma5'],
                        'ma10': indicators['ma10'],
                        'ma89': indicators['ma89'],
                        'ma120': indicators['ma120'],
                        'ma181': indicators['ma181'],
                        'ma250': indicators['ma250'],
                        'ema14': indicators['ema14'],
                        'ema_distance': indicators['ema_distance']
                    })

                    print(f"✅ 发现: {name}({symbol}) - 价格: {indicators['close']:.2f}")
                    print(
                        f"   均线系统: 5日{indicators['ma5']:.2f} > 10日{indicators['ma10']:.2f} > 89日{indicators['ma89']:.2f} > 120日{indicators['ma120']:.2f} > 181日{indicators['ma181']:.2f} > 250日{indicators['ma250']:.2f}")
                    print(f"   14日EMA: {indicators['ema14']:.2f}, 距离: {indicators['ema_distance']:.2f}%")

            except Exception as e:
                continue

            time.sleep(0.2)

        print(f"\n扫描完成，找到 {len(selected_stocks)} 只符合条件的股票")
        return selected_stocks


# 交易记录类
class TradeRecord:
    def __init__(self):
        self.trades = []

    def add_trade(self, buy_date, buy_price, sell_date, sell_price, profit_pct, reason):
        self.trades.append({
            'buy_date': buy_date,
            'buy_price': buy_price,
            'sell_date': sell_date,
            'sell_price': sell_price,
            'profit_pct': profit_pct,
            'reason': reason,
            'duration': (sell_date - buy_date).days if sell_date else None
        })

    def get_all_trades(self):
        return self.trades

    def get_total_trades(self):
        return len(self.trades)

    def get_win_trades(self):
        return sum(1 for trade in self.trades if trade['profit_pct'] > 0)

    def get_win_rate(self):
        if self.get_total_trades() == 0:
            return 0
        return (self.get_win_trades() / self.get_total_trades()) * 100


# 量化策略实现
class MultiMATrendStrategy(bt.Strategy):
    params = (
        ('printlog', True),
        ('stop_loss', 0.05),  # 5%止损
        ('open_loss', 0.04),  # 开盘跌幅4%
        ('close_loss', 0.05),  # 收盘跌幅5%
    )

    def __init__(self):
        # 计算移动平均线
        self.ma5 = bt.indicators.SimpleMovingAverage(self.data.close, period=5)
        self.ma10 = bt.indicators.SimpleMovingAverage(self.data.close, period=10)
        self.ma89 = bt.indicators.SimpleMovingAverage(self.data.close, period=89)
        self.ma120 = bt.indicators.SimpleMovingAverage(self.data.close, period=120)
        self.ma181 = bt.indicators.SimpleMovingAverage(self.data.close, period=181)
        self.ma250 = bt.indicators.SimpleMovingAverage(self.data.close, period=250)

        # 计算14日指数移动平均线
        self.ema14 = bt.indicators.ExponentialMovingAverage(self.data.close, period=14)

        # 跟踪变量
        self.order = None
        self.buyprice = None
        self.buycomm = None
        self.trade_records = TradeRecord()
        self.current_buy_date = None
        self.current_buy_price = None
        self.ema_break = False
        self.ema_break_date = None

    def next(self):
        if self.order is not None:
            return

        # 确保有足够的数据
        if len(self.data) < 250:
            return

        current_price = self.data.close[0]
        current_open = self.data.open[0]
        current_date = self.data.datetime.date(0)
        current_ema = self.ema14[0]

        # 卖出检查
        if self.position:
            # 5%止损
            if current_price <= self.buyprice * (1 - self.params.stop_loss):
                self.sell_stock('5%止损', current_date, current_price)
                return

            # 开盘跌幅4%
            if len(self.data) > 1 and current_open <= self.data.close[-1] * (1 - self.params.open_loss):
                self.sell_stock('开盘跌幅4%', current_date, current_price)
                return

            # EMA跌破处理
            if current_price < current_ema:
                if not self.ema_break:
                    # 第一次跌破EMA，标记并等待第二天
                    self.ema_break = True
                    self.ema_break_date = current_date
                else:
                    # 已经跌破过，检查是否连续两天在EMA下方
                    if current_date > self.ema_break_date:
                        self.sell_stock('EMA跌破不回', current_date, current_price)
                        return
            else:
                # 价格回到EMA上方，重置标记
                self.ema_break = False

        # 买入条件检查
        if not self.position:
            # 确保有足够的数据计算条件
            if len(self.ma5) < 2 or len(self.ema14) < 1:
                return

            # 1. 均线多头排列：5日>10日>89日>120日>181日>250日
            ma_condition = (self.ma5[0] > self.ma10[0] and
                            self.ma10[0] > self.ma89[0] and
                            self.ma89[0] > self.ma120[0] and
                            self.ma120[0] > self.ma181[0] and
                            self.ma181[0] > self.ma250[0] and
                            self.ma5[0] > self.ma5[-1])

            # 2. K线沿EMA且不破
            ema_condition = (current_price > self.ema14[0] and
                             abs(current_price - self.ema14[0]) / self.ema14[0] < 0.05)

            if ma_condition and ema_condition:
                self.buy_stock(current_date, current_price)
                self.ema_break = False

    def buy_stock(self, buy_date, buy_price):
        cash_available = self.broker.getcash()
        shares = int(cash_available / buy_price / 100) * 100
        if shares >= 100:
            self.order = self.buy(size=shares)
            self.current_buy_date = buy_date
            self.current_buy_price = buy_price
            if self.params.printlog:
                print(f"{buy_date}: 买入 - 价格: {buy_price:.2f}")

    def sell_stock(self, reason, sell_date, sell_price):
        self.order = self.close()
        if self.current_buy_date and self.current_buy_price:
            profit_pct = (sell_price / self.current_buy_price - 1) * 100
            self.trade_records.add_trade(
                self.current_buy_date, self.current_buy_price,
                sell_date, sell_price, profit_pct, reason
            )
            if self.params.printlog:
                print(f"{sell_date}: {reason} - 价格: {sell_price:.2f}, 盈亏: {profit_pct:+.2f}%")

            self.current_buy_date = None
            self.current_buy_price = None
            self.ema_break = False

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:
            if order.isbuy():
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            else:
                if self.current_buy_date and self.current_buy_price:
                    profit_pct = (order.executed.price / self.current_buy_price - 1) * 100
                    self.trade_records.add_trade(
                        self.current_buy_date, self.current_buy_price,
                        self.data.datetime.date(0), order.executed.price,
                        profit_pct, '订单完成'
                    )

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            pass

        self.order = None

    def stop(self):
        """策略结束时的处理"""
        if self.position:
            self.order = self.close()
            if self.current_buy_date and self.current_buy_price:
                profit_pct = (self.data.close[0] / self.current_buy_price - 1) * 100
                self.trade_records.add_trade(
                    self.current_buy_date, self.current_buy_price,
                    self.data.datetime.date(0), self.data.close[0],
                    profit_pct, '策略结束平仓'
                )


def get_stock_data(symbol, start_date, end_date):
    """获取股票数据"""
    try:
        df = ak.stock_zh_a_hist(symbol=symbol, period="daily",
                                start_date=start_date, end_date=end_date,
                                adjust="qfq")
        if df is None or df.empty:
            return None

        df = df.rename(columns={
            '日期': 'date', '开盘': 'open', '最高': 'high', '最低': 'low',
            '收盘': 'close', '成交量': 'volume'
        })
        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)
        df = df[['open', 'high', 'low', 'close', 'volume']]
        df = df.astype(float)

        # 确保数据足够
        if len(df) < 250:
            return None

        return df
    except:
        return None


def run_backtest(symbol, name, start_date, end_date, initial_cash=100000):
    """运行回测"""
    df = get_stock_data(symbol, start_date, end_date)
    if df is None or len(df) < 250:
        print(f"数据不足，跳过 {name}({symbol})")
        return None

    cerebro = bt.Cerebro()
    cerebro.broker.setcash(initial_cash)
    cerebro.broker.setcommission(commission=0.0003, stocklike=True)
    cerebro.addstrategy(MultiMATrendStrategy)

    data_feed = bt.feeds.PandasData(dataname=df, open=0, high=1, low=2, close=3, volume=4)
    cerebro.adddata(data_feed)

    print(f"回测 {name}({symbol})...")

    try:
        results = cerebro.run()
        strat = results[0]
    except Exception as e:
        print(f"回测 {name} 出错: {e}")
        return None

    final_value = cerebro.broker.getvalue()
    return_rate = ((final_value / initial_cash) - 1) * 100

    # 获取交易记录
    trade_records = strat.trade_records.get_all_trades()
    total_trades = len(trade_records)
    win_trades = sum(1 for trade in trade_records if trade['profit_pct'] > 0)
    win_rate = (win_trades / total_trades * 100) if total_trades > 0 else 0

    return {
        'symbol': symbol,
        'name': name,
        'trades': total_trades,
        'win_rate': win_rate,
        'return_rate': return_rate,
        'final_value': final_value,
        'trade_records': trade_records
    }


def print_trade_details(trade_records):
    """打印交易详情"""
    if not trade_records:
        print("    无交易记录")
        return

    print("    📋 交易详情:")
    for i, trade in enumerate(trade_records, 1):
        print(f"      {i:2d}. 买入: {trade['buy_date']} {trade['buy_price']:7.2f}元 | "
              f"卖出: {trade['sell_date']} {trade['sell_price']:7.2f}元 | "
              f"盈亏: {trade['profit_pct']:6.1f}% | "
              f"持仓: {trade['duration']:3}天 | "
              f"原因: {trade['reason']}")


def main():
    """主函数"""
    print("=" * 80)
    print("量化策略回测系统 - 多均线趋势策略（无520日均线）")
    print("=" * 80)
    print("策略规则:")
    print("1. 5/10/89/120/181/250日均线多头排列")
    print("2. K线沿14日EMA不破")
    print("3. 跌破EMA且第二天不回上方卖出")
    print("4. 开盘跌4%或收盘跌5%卖出")
    print("5. 5%止损")
    print("=" * 80)

    # 参数设置
    start_date = '20250101'
    end_date = '20251231'
    initial_cash = 100000
    max_stocks = 5000

    # 扫描股票
    scanner = StockScanner()
    selected_stocks = scanner.scan_stocks(start_date, end_date, max_stocks)

    if not selected_stocks:
        print("未找到满足条件的股票，使用默认股票测试")
        selected_stocks = [
            {'symbol': '000001', 'name': '平安银行'},
            {'symbol': '600036', 'name': '招商银行'}
        ]

    print(f"\n开始回测 {len(selected_stocks)} 只股票...")

    # 运行回测
    results = []
    for stock_info in selected_stocks:
        result = run_backtest(
            stock_info['symbol'],
            stock_info['name'],
            start_date,
            end_date,
            initial_cash
        )
        if result:
            results.append(result)
            print(f"{result['name']}({result['symbol']}): "
                  f"{result['trades']}次交易, "
                  f"胜率{result['win_rate']:.1f}%, "
                  f"收益{result['return_rate']:.1f}%")
        time.sleep(1)

    # 显示结果
    print("\n" + "=" * 80)
    print("回测结果汇总")
    print("=" * 80)

    if results:
        for i, result in enumerate(results, 1):
            print(f"{i:2d}. {result['name']}({result['symbol']}):")
            print(f"    交易次数: {result['trades']:3d}, 胜率: {result['win_rate']:5.1f}%")
            print(f"    收益率: {result['return_rate']:6.1f}%")
            print_trade_details(result.get('trade_records', []))
            print()

        # 统计信息
        total_return = sum(r['return_rate'] for r in results) / len(results)
        total_win_rate = sum(r['win_rate'] for r in results) / len(results)
        total_trades = sum(r['trades'] for r in results)

        print(f"📊 统计: 平均收益{total_return:.1f}%, 平均胜率{total_win_rate:.1f}%, 总交易次数{total_trades}")

    else:
        print("没有成功的回测结果")


if __name__ == '__main__':
    # 安装依赖: pip install backtrader pandas numpy matplotlib akshare tqdm
    main()